import { any } from "../../../core/logic";
import { Compute } from "./Compute";
import { taaCopy } from "./shaders/taaCopy";
import { taa, taa_multisampled } from "./shaders/taa";
import { Pass } from "../Pass";
import { Engine } from "../../../Engine";
import { Matrix4 } from "../../../core/math/Matrix4";
import { guid } from "../../../core/utils";

class TAAComputer extends Compute {
  constructor(opts = {}) {
    super(opts);
    this.initBuffer();
    // this.initPipeLine();
    this.blendFactor = any(opts.blendFactor, 0.3);
    this.jitterFrameIndex = 0;
  }
  set blendFactor(v) {
    this._blendFactor = v;
    // Engine.instance.queue.writeBuffer(this.taaSetting, 72, new Float32Array([v]));
    this.update();
    return this;
  }
  get blendFactor() {
    return this._blendFactor;
  }

  initBuffer() {
    this.preProjMatrix = new Matrix4();
    this.preViewMatrix = new Matrix4();
    let pixels = Engine.instance.size[0] * Engine.instance.size[1];
    this.taaSetting = Engine.instance.device.createBuffer({
      label: "taaSetting",
      size: 144,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
    });
    this.preColorBuffer = Engine.instance.device.createBuffer({
      label: "preColorBuffer",
      size: pixels * 4 * 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
    });
    this.preColorBuffer.id = guid();
    this.preColorTex = Engine.instance.device.createTexture({
      label: "preColorTex",
      size: Engine.instance.size,
      format: "rgba16float",
      mipLevelCount: 1,
      sampleCount: 1,
      usage:
        GPUTextureUsage.STORAGE_BINDING |
        GPUTextureUsage.COPY_DST |
        GPUTextureUsage.COPY_SRC |
        GPUTextureUsage.TEXTURE_BINDING,
    });
    this.preColorTex.id = guid();
    this.preColorTex.view = this.preColorTex.createView({ label: "preColorTex" })
    this.taaTexture = Engine.instance.device.createTexture({
      label: "taaTexture",
      size: Engine.instance.size,
      format: "rgba16float",
      mipLevelCount: 1,
      sampleCount: 1,
      usage:
        GPUTextureUsage.STORAGE_BINDING |
        GPUTextureUsage.COPY_DST |
        GPUTextureUsage.COPY_SRC |
        GPUTextureUsage.TEXTURE_BINDING,
    });
    this.taaTexture.id = guid();
    this.taaTexture.view = this.taaTexture.createView({ label: "taaTexture" })
    // this.sharpTaaTex = Engine.instance.device.createTexture({
    //   label: "sharpTaaTex",
    //   size: Engine.instance.size,
    //   format: "rgba16float",
    //   mipLevelCount: 1,
    //   sampleCount: 1,
    //   usage:
    //     GPUTextureUsage.STORAGE_BINDING |
    //     GPUTextureUsage.COPY_DST |
    //     GPUTextureUsage.COPY_SRC |
    //     GPUTextureUsage.TEXTURE_BINDING,
    // });
  }

  initPipeLine() {
    if (!this.copyComputePipeline) {
      this.copyComputePipeline = Engine.instance.device.createComputePipeline({
        layout: `auto`,
        compute: {
          module: Engine.instance.device.createShaderModule({
            code: taaCopy,
          }),
          entryPoint: "main",
        },
      });
    }
    if (!this.taaComputePipeline) {
      this.taaComputePipeline = Engine.instance.device.createComputePipeline({
        layout: `auto`,
        compute: {
          module: Engine.instance.device.createShaderModule({
            code: taa,
          }),
          entryPoint: "main",
        },
      });
      this.preColorTexSampler = Engine.instance.device.createSampler();
      this.colorTexSampler = Engine.instance.device.createSampler();
    }

  }
  initBindGroup() {
    if (
      this?.preColorTex?.width !== Engine.instance.size[0] ||
      this?.preColorTex?.height !== Engine.instance.size[1]
    ) {
      this.preColorBuffer.destroy();
      this.preColorBuffer = null;
      this.preColorTex.destroy();
      this.preColorTex = null;
      this.taaTexture.destroy();
      this.taaTexture = null;
      const pixels = Engine.instance.size[0] * Engine.instance.size[1];
      this.preColorBuffer = Engine.instance.device.createBuffer({
        label: "preColorBuffer",
        size: pixels * 4 * 4,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
      });
      this.preColorBuffer.id = guid();
      this.preColorTex = Engine.instance.device.createTexture({
        label: "preColorTex",
        size: Engine.instance.size,
        format: "rgba16float",
        mipLevelCount: 1,
        sampleCount: 1,
        usage:
          GPUTextureUsage.STORAGE_BINDING |
          GPUTextureUsage.COPY_DST |
          GPUTextureUsage.COPY_SRC |
          GPUTextureUsage.TEXTURE_BINDING,
      });
      this.preColorTex.id = guid();
      this.preColorTex.view = this.preColorTex.createView({ label: "preColorTex" })
      this.taaTexture = Engine.instance.device.createTexture({
        label: "taaTexture",
        size: Engine.instance.size,
        format: "rgba16float",
        mipLevelCount: 1,
        sampleCount: 1,
        usage:
          GPUTextureUsage.STORAGE_BINDING |
          GPUTextureUsage.COPY_DST |
          GPUTextureUsage.COPY_SRC |
          GPUTextureUsage.TEXTURE_BINDING,
      });
      this.taaTexture.id = guid();
      this.taaTexture.view = this.taaTexture.createView({ label: "taaTexture" })
      // this.sharpTaaTex = Engine.instance.device.createTexture({
      //   label: "sharpTaaTex",
      //   size: Engine.instance.size,
      //   format: "rgba16float",
      //   mipLevelCount: 1,
      //   sampleCount: 1,
      //   usage:
      //     GPUTextureUsage.STORAGE_BINDING |
      //     GPUTextureUsage.COPY_DST |
      //     GPUTextureUsage.COPY_SRC |
      //     GPUTextureUsage.TEXTURE_BINDING,
      // });
    }
    if (this.copyBindGroup && this.copyBindGroup.id != this.preColorBuffer.id + this.preColorTex.id) {
      this.copyBindGroup = null;
    }
    if (!this.copyBindGroup) {
      this.copyBindGroup = Engine.instance.device.createBindGroup({
        layout: this.copyComputePipeline.getBindGroupLayout(0),
        entries: [
          {
            binding: 0,
            resource: {
              buffer: this.preColorBuffer,
            },
          },
          {
            binding: 1,
            resource: this.preColorTex.view,
          },
        ],
      });
      this.copyBindGroup.id = this.preColorBuffer.id + this.preColorTex.id;
    }
    if (this.taaBindGroup && this.taaBindGroup.id != this.preColorBuffer.id + this.preColorTex.id + Pass.instance().readAttach.id + Pass.instance().postionTexture.id + this.taaTexture.id) {
      this.taaBindGroup = null;
    }
    if (!this.taaBindGroup) {
      const color = Pass.instance().readAttach.view
      this.taaBindGroup = Engine.instance.device.createBindGroup({
        label: "taaBindGroup",
        layout: this.taaComputePipeline.getBindGroupLayout(0),
        entries: [
          {
            binding: 1,
            resource: {
              buffer: this.taaSetting,
            },
          },
          {
            binding: 2,
            resource: {
              buffer: this.preColorBuffer,
            },
          },
          {
            binding: 3,
            resource: this.preColorTexSampler,
          },
          {
            binding: 4,
            resource: this.preColorTex.view,
          },
          {
            binding: 5,
            resource: Pass.instance().postionTexture.view,
          },
          {
            binding: 6,
            resource: this.colorTexSampler,
          },
          {
            binding: 7,
            resource: color,
          },
          {
            binding: 8,
            resource: this.taaTexture.view,
          },
        ],
      });
      this.taaBindGroup.id = this.preColorBuffer.id + this.preColorTex.id + Pass.instance().readAttach.id + Pass.instance().postionTexture.id + this.taaTexture.id;
    }

  }
  onCompute(commandEncoder) {
    this.initPipeLine()
    this.jitterFrameIndex++;
    this.initBindGroup();
    const camera = Engine.instance.scene.getCamera();
    let viewData = new Float32Array(this.preViewMatrix.elements);
    let projectData = new Float32Array(this.preProjMatrix.elements);
    let taaData = new Float32Array([
      ...projectData,
      ...viewData,
      this.jitterFrameIndex,
      this.blendFactor,
    ]);
    Engine.instance.queue.writeBuffer(this.taaSetting, 0, taaData);
    const workerSizeX = Math.ceil(Engine.instance.size[0] / 8);
    const workerSizeY = Math.ceil(Engine.instance.size[1] / 8);
    const workerSizeZ = 1;
    const computePass = commandEncoder.beginComputePass();

    // taa
    computePass.setPipeline(this.taaComputePipeline);
    computePass.setBindGroup(0, this.taaBindGroup);
    computePass.dispatchWorkgroups(workerSizeX, workerSizeY, workerSizeZ);

    // copy
    computePass.setPipeline(this.copyComputePipeline);
    computePass.setBindGroup(0, this.copyBindGroup);
    computePass.dispatchWorkgroups(workerSizeX, workerSizeY, workerSizeZ);
    //
    computePass.end();
    this.preProjMatrix.copy(camera.project);
    this.preViewMatrix.copy(camera.view);
    // this.onCompute1(commandEncoder)
  }
}
export { TAAComputer };
